package it.uniroma2.exp.dtk;

import it.uniroma2.dtk.dt.DT;
import it.uniroma2.dtk.dt.GenericDT;
import it.uniroma2.dtk.dt.route.GenericRouteDT;
import it.uniroma2.dtk.dt.subpath.GenericSubPathDT;
import it.uniroma2.exp.AbstractExperiment;
import it.uniroma2.exp.tools.AvgVarCalculator;
import it.uniroma2.tk.RouteTreeKernel;
import it.uniroma2.tk.SubpathTreeKernel;
import it.uniroma2.tk.TreeKernel;
import it.uniroma2.util.math.ArrayMath;
import it.uniroma2.util.math.statistical.SpearmanCorrelation;
import it.uniroma2.util.tree.RandomTreeGenerator;
import it.uniroma2.util.tree.Tree;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileReader;
import java.io.PrintStream;
import java.util.Vector;

/**
 * @author Lorenzo Dell'Arciprete
 * 
 * This experiment evaluates how similar are two kernel results for trees taken from a certain corpus.
 * The available corpora are artificially generated trees, question classification trees and
 * textual entrailment recognition trees [not available yet].
 * Other than the Tree Kernel and Distributed Tree Kernel, a Zero Tree Kernel is included, 
 * that is simply a Tree Kernel with lambda = 0.
 * The kernel value similarity measures are Spearman Correlation and Ratio 
 */
public class KernelComparison extends AbstractExperiment {
	
	public enum OPS { TK, DTK, ZTK };
	public enum TYPES { artificial, rte, qc };
	//The location of the Question Classification trees file (in SVM Light format)
	public final String inputTreesQC = "/home/lorenzo/esperimenti/dtk/qc/correlationTestData/corr_test_data.dat";
	
	private static final String rte_base_dir = "/home/lorenzo/workspace/arte/data/experiments/preprocessing/CPW/";

	public static final String[] inputTreePairsRTE = {
		rte_base_dir  + "RTE1_dev.xml"
		,rte_base_dir  + "RTE1_test.xml"
		,rte_base_dir  + "RTE2_dev.xml"
		,rte_base_dir  + "RTE2_test.xml"
		,rte_base_dir  + "RTE3_dev.xml"
		,rte_base_dir  + "RTE3_test.xml"
		,rte_base_dir  + "RTE4_test.xml"
		,rte_base_dir  + "RTE5_dev.xml"
		,rte_base_dir  + "RTE5_test.xml"
	};

	private DT dtk = null;
	private int kernel = 0;
	
	//The two kernels to be compared
	private OPS OP1;
	private OPS OP2;
	//The corpora for the trees
	private TYPES TYPE;
	
	public static void main(String [] argv) throws Exception {
		KernelComparison kc = new KernelComparison();
		try {
			kc.setOutputStream(new PrintStream(new File("kernel_comparison.dat")));
		} catch (FileNotFoundException e) {
			e.printStackTrace();
			System.exit(0);
		}
		kc.setVectorSizeArray(new int[] {1024, 2048, 4096, 8192});
		kc.setRandomTreeNodes(30);
		kc.setLambdaArray(new double[]{0.2, 0.4, 0.6, 0.8, 1});
		//The corpus type, the operations pair and the kernel type
		//Custom parameters are:
		//Corpus type: 0=artificial, 1=qc, 2=rte
		//Operations pair: 0=DTKvsTK, 1=DTKvsZTK, 2=ZTKvsTK
		//Kernel type: 0=Tree Kernel, 1=Subpath Kernel, 2=Route Kernel
		kc.setCustomParameters(new int[][] {{0,1,2},{0},{0,1,2}});
		kc.runAll();
	}

	@Override
	protected void runExperiment() throws Exception {
		TYPE = TYPES.values()[currentCustomParameters[0]];
		int opPair = currentCustomParameters[1];
		if (opPair == 0) {
			OP1 = OPS.DTK;
			OP2 = OPS.TK;
		}
		else if (opPair == 1) {
			OP1 = OPS.DTK;
			OP2 = OPS.ZTK;
		}
		else if (opPair == 2) {
			OP1 = OPS.ZTK;
			OP2 = OPS.TK;
		} 
		kernel = currentCustomParameters[2];
		if (kernel == 0) {
			TreeKernel.lexicalized = lexicalized;
			TreeKernel.lambda = lambda;
		}
		else if (kernel == 1)
			SubpathTreeKernel.lambda = lambda;
		else if (kernel == 2)
			RouteTreeKernel.lambda = lambda;
		usePos = TYPE != TYPES.artificial && kernel != 1;
		if (kernel == 0)
			dtk = new GenericDT(randomOffset, vectorSize, usePos, lexicalized, lambda, compositionType);
		else if (kernel == 1)
			dtk = new GenericSubPathDT(randomOffset, vectorSize, usePos, lambda, compositionType);
		else if (kernel == 2)
			dtk = new GenericRouteDT(randomOffset, vectorSize, false, lambda, compositionType);
		double spearmanCorrelation = 0;
		AvgVarCalculator avgRatioCor = new AvgVarCalculator();
		switch(TYPE) {
		case artificial: spearmanCorrelation = computeArtificial(avgRatioCor); break;
		case rte: spearmanCorrelation = computeRTE(avgRatioCor); break;
		case qc: spearmanCorrelation = computeQC(avgRatioCor); break;
		}
		out.println("Spearman\tRatio");
		out.print(String.format("%.3f\t", spearmanCorrelation));
		out.println(String.format("%.3f\t%.3f\t", avgRatioCor.getAvg(), avgRatioCor.getVar()));
	}
	
	public double computeRTE(AvgVarCalculator ratio) throws Exception {
		double[] resultsSpear = new double[inputTreePairsRTE.length];
		for (int n=0; n<inputTreePairsRTE.length; n++) {
			BufferedReader in = new BufferedReader(new FileReader(inputTreePairsRTE[n]));
			SpearmanCorrelation corr = new SpearmanCorrelation();
			Vector<Double> a_vec = new Vector<Double>(); 
			Vector<Double> b_vec = new Vector<Double>();
			int count = 0;
			while (true) {
				String treet = in.readLine();
				if (treet == null)
					break;
				else
					treet = treet.trim();
				if (!treet.startsWith("<t>"))
					continue;
				String treeh = in.readLine().trim();
				if (!treeh.startsWith("<h>"))
					continue;
				treet = treet.substring(3, treet.indexOf("</t>")).trim(); 
				Tree xt = Tree.fromPennTree(treet);
				treeh = treeh.substring(3, treeh.indexOf("</h>")).trim(); 
				Tree xh = Tree.fromPennTree(treeh);
				double value1 = op(xt, xh, OP1);
				double value2 = op(xt, xh, OP2);
				a_vec.add(value1);
				b_vec.add(value2);
				if (value2 != 0)
					ratio.addSample(value1/value2);
				out.print('.');
				if (count%100 == 0) out.println("" +count);
				count += 2;
			}
			float[] a = new float[a_vec.size()];
			for (int i=0; i< a_vec.size();i++) a[i]= a_vec.elementAt(i).floatValue();
			float[] b = new float[b_vec.size()];
			for (int i=0; i< b_vec.size();i++) b[i]= b_vec.elementAt(i).floatValue();
			resultsSpear[n] = corr.spearmanCorrelationCoefficient(a, b);
			in.close();
			out.println("end");
		}
		double tot = 0;
		for (double d : resultsSpear)
			tot += d;
		return tot/inputTreePairsRTE.length;
	}
	
	public double computeArtificial(AvgVarCalculator ratio) throws Exception {
		RandomTreeGenerator rtg = new RandomTreeGenerator(randomOffset);
		Vector<Double> a_vec = new Vector<Double>(); 
		Vector<Double> b_vec = new Vector<Double>();
		SpearmanCorrelation corr = new SpearmanCorrelation();
		for (int i=0; i<1000; i++) {
			out.print('.');
			if (i%100 == 0) out.println("" +i);
			String treeh = rtg.generateRandomTree(randomTreeLabels, randomTreeMaxDegree, randomTreeNodes);
			Tree xh = Tree.fromPennTree(treeh);
			String treet = rtg.generateRandomTree(randomTreeLabels, randomTreeMaxDegree, randomTreeNodes);
			Tree xt = Tree.fromPennTree(treet);
			try {
				double value1 = op(xt, xh, OP1);
				double value2 = op(xt, xh, OP2);
				a_vec.add(value1);
				b_vec.add(value2);
				if (value2 != 0)
					ratio.addSample(value1/value2);
			} catch (Exception e) {
				out.println(e.getMessage());
			}
		}
		out.println("end");
		float[] a = new float[a_vec.size()];
		for (int i=0; i< a_vec.size();i++) a[i]= a_vec.elementAt(i).floatValue();
		float[] b = new float[b_vec.size()];
		for (int i=0; i< b_vec.size();i++) b[i]= b_vec.elementAt(i).floatValue();
		return corr.spearmanCorrelationCoefficient(a, b);
	}
	
	public double computeQC(AvgVarCalculator ratio) throws Exception {
		Vector<Double> a_vec = new Vector<Double>(); 
		Vector<Double> b_vec = new Vector<Double>();
		SpearmanCorrelation corr = new SpearmanCorrelation();
		BufferedReader in = new BufferedReader(new FileReader(inputTreesQC));
		int count = 0;
		while (true) {
			String treet = in.readLine();
			String treeh = in.readLine();
			if (treet == null || treeh == null)
				break;
			treet = treet.substring(treet.indexOf("|BT|")+4, treet.indexOf("|ET|")).trim(); 
			Tree xt = Tree.fromPennTree(treet);
			treeh = treeh.substring(treeh.indexOf("|BT|")+4, treeh.indexOf("|ET|")).trim(); 
			Tree xh = Tree.fromPennTree(treeh);
			double value1 = op(xt, xh, OP1);
			double value2 = op(xt, xh, OP2);
			a_vec.add(value1);
			b_vec.add(value2);
			if (value2 != 0)
				ratio.addSample(value1/value2);
			out.print('.');
			if (count%100 == 0) out.println("" +count);
			count += 2;
		}
		in.close();
		out.println("end");
		float[] a = new float[a_vec.size()];
		for (int i=0; i< a_vec.size();i++) a[i]= a_vec.elementAt(i).floatValue();
		float[] b = new float[b_vec.size()];
		for (int i=0; i< b_vec.size();i++) b[i]= b_vec.elementAt(i).floatValue();
		return corr.spearmanCorrelationCoefficient(a, b);
	}
	
	public double op(Tree xt, Tree xh, OPS op) throws Exception {
		switch(op) {
			case TK: return TK(xt, xh); 
			case DTK: return DTK(xt, xh); 
			case ZTK: return ZTK(xt, xh); 
			default: throw new Exception("Unknown operation!");
		}
	}
	
	public double TK(Tree xt, Tree xh) {
		if (kernel == 0)
			return TreeKernel.value(xt, xh);
		else if (kernel == 1)
			return SubpathTreeKernel.value(xt, xh);
		else if (kernel == 2)
			return RouteTreeKernel.value(xt, xh);
		else return 0;
	}
	
	public double DTK(Tree xt, Tree xh) throws Exception {
		double[] vectort = dtk.dt(xt);
		double[] vectorh = dtk.dt(xh);
		return ArrayMath.dot(vectort,vectorh);
	}
	
	public double ZTK(Tree xt, Tree xh) {
		double result = 0;
		if (kernel == 0) {
			double oldLambda = TreeKernel.lambda;
			TreeKernel.lambda = 0;
			result = TK(xt, xh);
			TreeKernel.lambda = oldLambda;
		}
		else if (kernel == 1) {
			double oldLambda = SubpathTreeKernel.lambda;
			SubpathTreeKernel.lambda = 0;
			result = TK(xt, xh);
			SubpathTreeKernel.lambda = oldLambda;
		}
		else if (kernel == 2) {
			double oldLambda = RouteTreeKernel.lambda;
			RouteTreeKernel.lambda = 0;
			result = TK(xt, xh);
			RouteTreeKernel.lambda = oldLambda;
		}
		return result;
	}
}
